#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
智能AI助手API - 集成LangChain和ChatGPT
"""

import os
import sqlite3
from typing import Any, Optional
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
from langchain_openai import ChatOpenAI
from langchain.schema import HumanMessage, SystemMessage
from langchain.tools import Tool
from langchain.agents import initialize_agent, AgentType
import json
from datetime import datetime

router = APIRouter()

# 对话记忆存储（实际项目中应该使用数据库）
conversation_memory = {}

# OpenAI配置
OPENAI_API_KEY = 'sk-W3u9YKJyZf1yS4G6Lv1ynfe299eoB0u2RDg9FOJC3hkvxf9e'
BASE_URL = "https://api.chatanywhere.tech"
MODEL = "deepseek-v3"


class QueryRequest(BaseModel):
    query: str
    context: str = ""
    conversation_id: str = ""  # 对话ID，用于记忆功能


class QueryResponse(BaseModel):
    answer: str
    data: Any = []
    sql_query: str = ""
    error: str = ""
    conversation_id: str = ""  # 返回对话ID


def get_db_connection():
    """获取SQLite数据库连接"""
    db_path = os.path.join(os.path.dirname(__file__), "..", "database", "shop.db")
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    return conn


def get_sales_info(query: str) -> str:
    """获取销售信息"""
    try:
        conn = get_db_connection()
        cursor = conn.cursor()

        # 获取销售概览
        cursor.execute("""
            SELECT 
                COUNT(*) as total_orders,
                SUM(price * quantity) as total_sales,
                SUM(quantity) as total_quantity,
                AVG(price * quantity) as avg_order_value
            FROM orders
        """)

        row = cursor.fetchone()
        total_orders = row[0] if row[0] else 0
        total_sales = row[1] if row[1] else 0
        total_quantity = row[2] if row[2] else 0
        avg_order_value = row[3] if row[3] else 0

        # 获取最近7天销售趋势
        cursor.execute("""
            SELECT 
                order_date,
                COUNT(*) as daily_orders,
                SUM(price * quantity) as daily_sales
            FROM orders
            GROUP BY order_date
            ORDER BY order_date DESC
            LIMIT 7
        """)

        daily_trend = []
        for row in cursor.fetchall():
            daily_trend.append({
                "date": row[0],
                "orders": row[1],
                "sales": row[2]
            })

        # 获取热销商品
        cursor.execute("""
            SELECT 
                p.name,
                SUM(o.quantity) as total_sold,
                SUM(o.price * o.quantity) as total_revenue
            FROM orders o
            JOIN products p ON o.product_id = p.id
            GROUP BY o.product_id, p.name
            ORDER BY total_sold DESC
            LIMIT 5
        """)

        top_products = []
        for row in cursor.fetchall():
            top_products.append({
                "name": row[0],
                "sold": row[1],
                "revenue": row[2]
            })

        conn.close()

        result = f"""📊 销售情况概览：

🎯 总体数据：
• 总订单数：{total_orders} 个
• 总销售额：¥{total_sales:.2f}
• 总销售商品数：{total_quantity} 件
• 平均订单金额：¥{avg_order_value:.2f}

📈 最近7天销售趋势：
"""

        for trend in daily_trend:
            result += f"• {trend['date']}: {trend['orders']} 单，¥{trend['sales']:.1f}\n"

        result += "\n🔥 热销商品TOP5：\n"
        for i, product in enumerate(top_products, 1):
            result += f"• {i}. {product['name']}: 销量{product['sold']}件，收入¥{product['revenue']:.1f}\n"

        return result

    except Exception as e:
        return f"获取销售信息时出现错误: {str(e)}"


def get_inventory_info(query: str) -> str:
    """获取库存信息"""
    try:
        conn = get_db_connection()
        cursor = conn.cursor()

        # 获取库存概览
        cursor.execute("""
            SELECT 
                COUNT(*) as total_products,
                SUM(quantity) as total_inventory,
                AVG(quantity) as avg_inventory,
                COUNT(CASE WHEN quantity <= 10 THEN 1 END) as low_stock
            FROM products
        """)

        row = cursor.fetchone()
        total_products = row[0] if row[0] else 0
        total_inventory = row[1] if row[1] else 0
        avg_inventory = row[2] if row[2] else 0
        low_stock = row[3] if row[3] else 0

        # 获取分类库存
        cursor.execute("""
            SELECT 
                c.name,
                COUNT(p.id) as product_count,
                SUM(p.quantity) as total_quantity
            FROM products p
            JOIN categories c ON p.category_id = c.id
            GROUP BY c.id, c.name
        """)

        category_inventory = []
        for row in cursor.fetchall():
            category_inventory.append({
                "category": row[0],
                "product_count": row[1],
                "total_quantity": row[2]
            })

        # 获取低库存商品
        cursor.execute("""
            SELECT 
                name,
                quantity
            FROM products
            WHERE quantity <= 10
            ORDER BY quantity ASC
            LIMIT 5
        """)

        low_stock_products = []
        for row in cursor.fetchall():
            low_stock_products.append({
                "name": row[0],
                "quantity": row[1]
            })

        conn.close()

        result = f"""📦 库存情况概览：

🎯 总体数据：
• 商品种类：{total_products} 种
• 总库存量：{total_inventory} 件
• 平均库存量：{avg_inventory:.1f} 件
• 低库存商品数（≤10）：{low_stock} 种

📋 分类库存：
"""

        for cat in category_inventory:
            result += f"• {cat['category']}: {cat['product_count']} 种商品，{cat['total_quantity']} 件\n"

        result += "\n⚠️ 低库存预警（TOP5）：\n"
        if low_stock_products:
            for i, product in enumerate(low_stock_products, 1):
                result += f"• {i}. {product['name']}: {product['quantity']} 件\n"
        else:
            result += "• 无低库存商品\n"

        return result

    except Exception as e:
        return f"获取库存信息时出现错误: {str(e)}"


def get_product_info(query: str) -> str:
    """获取产品信息"""
    try:
        conn = get_db_connection()
        cursor = conn.cursor()

        # 获取产品概览
        cursor.execute("""
            SELECT 
                COUNT(*) as total_products,
                AVG(price) as avg_price,
                MIN(price) as min_price,
                MAX(price) as max_price
            FROM products
        """)

        row = cursor.fetchone()
        total_products = row[0] if row[0] else 0
        avg_price = row[1] if row[1] else 0
        min_price = row[2] if row[2] else 0
        max_price = row[3] if row[3] else 0

        # 获取分类分布
        cursor.execute("""
            SELECT 
                c.name,
                COUNT(p.id) as product_count,
                AVG(p.price) as avg_price
            FROM products p
            JOIN categories c ON p.category_id = c.id
            GROUP BY c.id, c.name
        """)

        category_distribution = []
        for row in cursor.fetchall():
            category_distribution.append({
                "category": row[0],
                "product_count": row[1],
                "avg_price": row[2]
            })

        conn.close()

        result = f"""🛍️ 产品信息概览：

🎯 总体数据：
• 商品总数：{total_products} 种
• 平均价格：¥{avg_price:.2f}
• 最低价格：¥{min_price:.2f}
• 最高价格：¥{max_price:.2f}

📋 分类分布：
"""

        for cat in category_distribution:
            result += f"• {cat['category']}: {cat['product_count']} 种商品，平均价格 ¥{cat['avg_price']:.2f}\n"

        return result

    except Exception as e:
        return f"获取产品信息时出现错误: {str(e)}"


def get_customer_info(query: str) -> str:
    """获取客户信息"""
    try:
        conn = get_db_connection()
        cursor = conn.cursor()

        # 获取客户概览
        cursor.execute("""
            SELECT 
                COUNT(DISTINCT customer_email) as total_customers,
                SUM(price * quantity) as total_spending,
                AVG(price * quantity) as avg_spending
            FROM orders
        """)

        row = cursor.fetchone()
        total_customers = row[0] if row[0] else 0
        total_spending = row[1] if row[1] else 0
        avg_spending = row[2] if row[2] else 0

        # 获取消费排行
        cursor.execute("""
            SELECT 
                customer_name,
                COUNT(*) as order_count,
                SUM(price * quantity) as total_spending
            FROM orders
            GROUP BY customer_email, customer_name
            ORDER BY total_spending DESC
            LIMIT 5
        """)

        top_customers = []
        for row in cursor.fetchall():
            top_customers.append({
                "name": row[0],
                "order_count": row[1],
                "total_spending": row[2]
            })

        conn.close()

        result = f"""👥 客户信息概览：

🎯 总体数据：
• 客户总数：{total_customers} 人
• 总消费金额：¥{total_spending:.2f}
• 平均消费金额：¥{avg_spending:.2f}

🏆 消费排行TOP5：
"""

        for i, customer in enumerate(top_customers, 1):
            result += f"• {i}. {customer['name']}: {customer['order_count']} 单，消费 ¥{customer['total_spending']:.2f}\n"

        return result

    except Exception as e:
        return f"获取客户信息时出现错误: {str(e)}"


def get_category_info(query: str) -> str:
    """获取分类信息"""
    try:
        conn = get_db_connection()
        cursor = conn.cursor()

        # 获取分类概览
        cursor.execute("""
            SELECT 
                COUNT(*) as total_categories
            FROM categories
        """)

        total_categories = cursor.fetchone()[0] or 0

        # 获取分类详情
        cursor.execute("""
            SELECT 
                c.name,
                COUNT(p.id) as product_count,
                AVG(p.price) as avg_price,
                SUM(p.quantity) as total_quantity
            FROM categories c
            LEFT JOIN products p ON c.id = p.category_id
            GROUP BY c.id, c.name
        """)

        category_details = []
        for row in cursor.fetchall():
            category_details.append({
                "name": row[0],
                "product_count": row[1],
                "avg_price": row[2] or 0,
                "total_quantity": row[3] or 0
            })

        conn.close()

        result = f"""🏷️ 分类信息概览：

🎯 总体数据：
• 分类总数：{total_categories} 个

📋 分类详情：
"""

        for cat in category_details:
            result += f"• {cat['name']}: {cat['product_count']} 种商品，平均价格 ¥{cat['avg_price']:.2f}，总库存 {cat['total_quantity']} 件\n"

        return result

    except Exception as e:
        return f"获取分类信息时出现错误: {str(e)}"


def get_summary_info(query: str) -> str:
    """获取统计汇总"""
    try:
        conn = get_db_connection()
        cursor = conn.cursor()

        # 获取基础数据
        cursor.execute("SELECT COUNT(*) FROM products")
        total_products = cursor.fetchone()[0] or 0

        cursor.execute("SELECT COUNT(*) FROM orders")
        total_orders = cursor.fetchone()[0] or 0

        cursor.execute("SELECT COUNT(DISTINCT customer_email) FROM orders")
        unique_customers = cursor.fetchone()[0] or 0

        conn.close()

        result = f"""📈 系统统计汇总：

🏪 基础数据：
• 商品总数：{total_products} 种
• 订单总数：{total_orders} 单
• 独立客户：{unique_customers} 人
"""

        return result

    except Exception as e:
        return f"获取统计汇总时出现错误: {str(e)}"


# 创建LangChain工具
tools = [
    Tool(
        name="sales_info",
        func=get_sales_info,
        description="获取销售情况信息，包括订单数量、销售额、热销商品等"
    ),
    Tool(
        name="inventory_info",
        func=get_inventory_info,
        description="获取库存情况信息，包括商品数量、缺货预警、分类库存等"
    ),
    Tool(
        name="product_info",
        func=get_product_info,
        description="获取产品信息，包括产品概览、分类分布、价格分布等"
    ),
    Tool(
        name="customer_info",
        func=get_customer_info,
        description="获取客户信息，包括客户数量、消费排行、订单分布等（仅限管理员）"
    ),
    Tool(
        name="category_info",
        func=get_category_info,
        description="获取分类信息，包括分类统计、商品分布、价格分析等"
    ),
    Tool(
        name="summary_info",
        func=get_summary_info,
        description="获取统计汇总信息，包括整体概览、业务指标、趋势分析等"
    )
]

# 初始化ChatGPT模型
llm = ChatOpenAI(
    openai_api_key=OPENAI_API_KEY,
    base_url=BASE_URL,
    model=MODEL,
    temperature=0.7
)

# 系统提示词
system_prompt = """你是一个智能的购物管理系统AI助手，专门帮助管理员了解和分析业务数据。

你的主要职责：
1. 理解用户的各种问题（包括时间、天气、闲聊等）
2. 对于业务相关问题，必须使用相应的工具查询数据库获取真实数据
3. 提供友好、自然的对话体验
4. 回答要完整、准确，基于查询到的真实数据
5. 对于非业务问题，提供合适的回答
6. 如果上下文指示 is_admin=False，禁止使用 customer_info 工具，返回提示“抱歉，用户无权访问客户信息。”

可用工具：
- sales_info: 查询销售情况，包括订单数量、销售额、热销商品等
- inventory_info: 查询库存情况，包括商品数量、缺货预警、分类库存等
- product_info: 查询产品信息，包括产品概览、分类分布、价格分布等
- customer_info: 查询客户信息，包括客户数量、消费排行、订单分布等（仅限管理员）
- category_info: 查询分类信息，包括分类统计、商品分布、价格分析等
- summary_info: 查询统计汇总，包括整体概览、业务指标、趋势分析等

重要：当用户询问业务数据时，你必须使用相应的工具查询数据库，然后基于查询结果提供完整、准确的回答。不要只是描述工具，要真正执行查询并返回结果。如果 is_admin=False，禁止使用 customer_info 工具。"""

# 初始化Agent
agent = initialize_agent(
    tools,
    llm,
    agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True,
    handle_parsing_errors=True
)


@router.post("/ai-assistant/query")
async def query_ai_assistant(request: QueryRequest):
    """智能AI助手查询接口"""
    try:
        # 解析 context 中的 is_admin
        is_admin = False
        if request.context.startswith("is_admin="):
            is_admin = request.context.split("=")[1].lower() == "true"

        # 处理查询
        if not request.query.strip():
            return QueryResponse(
                answer="您好！我是智能AI助手，可以帮您了解购物管理系统的各种信息。请告诉我您想了解什么？",
                data=[],
                sql_query="",
                error="",
                conversation_id=""
            )

        # 如果查询涉及客户信息且 is_admin=False，拒绝查询
        if not is_admin and "客户信息" in request.query:
            return QueryResponse(
                answer="抱歉，用户无权访问客户信息。",
                data=[],
                sql_query="",
                error="权限不足",
                conversation_id=request.conversation_id or ""
            )

        # 生成或获取对话ID
        if not request.conversation_id:
            conversation_id = f"conv_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{len(conversation_memory)}"
        else:
            conversation_id = request.conversation_id

        # 获取对话历史
        conversation_history = conversation_memory.get(conversation_id, [])

        # 添加当前查询到历史
        conversation_history.append({
            "role": "user",
            "content": request.query,
            "timestamp": datetime.now().isoformat()
        })

        # 构建带上下文的查询
        context_query = f"系统提示: {system_prompt}\n\n"
        if conversation_history:
            context_query += "对话历史:\n"
            for msg in conversation_history[-3:]:  # 只保留最近3条消息
                context_query += f"{msg['role']}: {msg['content']}\n"
            context_query += "\n"
        context_query += f"当前用户问题: {request.query}\n\n请基于对话历史上下文，完整回答用户问题。"

        # 使用LangChain Agent处理查询
        result = agent.run(context_query)

        # 添加AI回答到历史
        conversation_history.append({
            "role": "assistant",
            "content": result,
            "timestamp": datetime.now().isoformat()
        })

        # 保存对话历史
        conversation_memory[conversation_id] = conversation_history

        return QueryResponse(
            answer=result,
            data=[],
            sql_query="",
            error="",
            conversation_id=conversation_id
        )

    except Exception as e:
        return QueryResponse(
            answer=f"抱歉，处理您的查询时出现了错误: {str(e)}",
            data=str(e),
            sql_query="",
            error=str(e),
            conversation_id=request.conversation_id or ""
        )